from fastapi import APIRouter, HTTPException, Depends
from pydantic import BaseModel

from auth_service.auths import verify_token
from retriever.services import baseline_search_chunks
from tables.knowledge import KG
from utils.main import str_to_hash

router = APIRouter(
    prefix="/rag/retrieve",
    tags=["retrieve management"],
    responses={404: {"description": "Not found"}}
)


class SearchInfo(BaseModel):
    query: str
    rag_id: str


@router.post("/search", summary="检索")
async def search_api(
        body: SearchInfo,
        current_user=Depends(verify_token)
):
    try:
        kg = KG.get_knowledge(userid=current_user.id, id=body.rag_id)
        if not kg.is_set:
            raise Exception("知识库未设置完成")
        if kg.rag_type == 1:
            res = await baseline_search_chunks(
                query=body.query,
                llm=kg.llm,
                embed=kg.embed,
                collection_name=str_to_hash(body.rag_id),
                retrieval_type=kg.retrieval_type,
                reranker=kg.reranker,
                content_augmentation=kg.content_augmentation,
                top_k=kg.top_k,
                alpha=kg.mix_threshold,
            )
            return {"code": 200, "data": res, "message": "获取检索结果成功"}
        else:
            raise HTTPException(status_code=400, detail="知识库类型错误")
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
